#include "imu_processor/imu_processor.hpp"
#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <tf2_ros/static_transform_broadcaster.h>
#include <tf2/LinearMath/Vector3.h>
#include <std_srvs/srv/trigger.hpp>
#include <memory>
#include <array>
#include <cmath>


using namespace std::chrono_literals;

SimpleKalmanFilter::SimpleKalmanFilter(double process_noise, double measurement_noise)
    : x_(0.0), P_(1.0), Q_(process_noise), R_(measurement_noise), initialized_(false) {}


double SimpleKalmanFilter::filter(double measurement) {
    if(!initialized_) {
        x_ = measurement;
        initialized_ = true;
        return x_;
    }

    //预测步骤
    double P_pred = P_ + Q_;

    //更新步骤
    double K = P_pred / (P_pred + R_);  //卡尔曼增益
    x_ = x_ + K * (measurement - x_);   //状态更新
    P_ = (1 - K) * P_pred;              //协方差更新

    return x_;
}

void SimpleKalmanFilter::update_process_noise(double new_process_noise) {
    Q_ = new_process_noise;
}

void SimpleKalmanFilter::reset() {
    initialized_ = false;
    x_ = 0.0;
    P_ = 1.0;
}

//IMUCalibrator 实现
void IMUCalibrator::start_calibration() {
    is_calibrating_ = true;
    sample_count_ = 0;
    gyro_x_samples.clear();
    gyro_y_samples.clear();
    gyro_z_samples.clear();
    accel_x_samples.clear();
    accel_y_samples.clear();
    accel_z_samples.clear();
}

void IMUCalibrator::update_calibration(const sensor_msgs::msg::Imu::SharedPtr imu_msg) {
    if (!is_calibrating_ || sample_count_ >= max_samples_) {
        return;
    }

    // 收集陀螺仪数据
    gyro_x_samples.push_back(imu_msg->angular_velocity.x);
    gyro_y_samples.push_back(imu_msg->angular_velocity.y);
    gyro_z_samples.push_back(imu_msg->angular_velocity.z);

    // 收集加速度计数据
    accel_x_samples.push_back(imu_msg->linear_acceleration.x);
    accel_y_samples.push_back(imu_msg->linear_acceleration.y);
    accel_z_samples.push_back(imu_msg->linear_acceleration.z);

    sample_count_++;
}

void IMUCalibrator::compute_bias() {
    if (sample_count_ == 0) return;

    // 计算陀螺仪零偏（静止时角速度应为0）
    bias_.gyro_x = std::accumulate(gyro_x_samples.begin(), gyro_x_samples.end(), 0.0) / sample_count_;
    bias_.gyro_y = std::accumulate(gyro_y_samples.begin(), gyro_y_samples.end(), 0.0) / sample_count_;
    bias_.gyro_z = std::accumulate(gyro_z_samples.begin(), gyro_z_samples.end(), 0.0) / sample_count_;

    // 计算加速度计零偏
    bias_.accel_x = std::accumulate(accel_x_samples.begin(), accel_x_samples.end(), 0.0) / sample_count_;
    bias_.accel_y = std::accumulate(accel_y_samples.begin(), accel_y_samples.end(), 0.0) / sample_count_;
    bias_.accel_z = std::accumulate(accel_z_samples.begin(), accel_z_samples.end(), 0.0) / sample_count_;
}

sensor_msgs::msg::Imu IMUCalibrator::apply_calibration(const sensor_msgs::msg::Imu& raw_imu) const {
    sensor_msgs::msg::Imu corrected_imu = raw_imu;
    
    // 应用陀螺仪零偏校正
    corrected_imu.angular_velocity.x -= bias_.gyro_x;
    corrected_imu.angular_velocity.y -= bias_.gyro_y;
    corrected_imu.angular_velocity.z -= bias_.gyro_z;
    
    // 应用加速度计零偏校正
    corrected_imu.linear_acceleration.x -= bias_.accel_x;
    corrected_imu.linear_acceleration.y -= bias_.accel_y;
    corrected_imu.linear_acceleration.z -= bias_.accel_z;
    
    return corrected_imu;
}


ImuProcessor::ImuProcessor() : Node("imu_processor")
{
    //参数声明
    this->declare_parameter("input_imu_topic", "/lidar_imu");
    this->declare_parameter("output_imu_topic", "/processed_imu");
    this->declare_parameter("imu_frame", "imu_link");
    this->declare_parameter("base_frame", "base_link");
    this->declare_parameter("publish_tf", true);
    this->declare_parameter("imu_position_x", 0.1);
    this->declare_parameter("imu_position_y", 0.0);
    this->declare_parameter("imu_position_z", 0.05);
    this->declare_parameter("imu_rotation_x", 0.0);
    this->declare_parameter("imu_rotation_y", 0.0);
    this->declare_parameter("imu_rotation_z", 0.0);
    this->declare_parameter("imu_rotation_w", 1.0);
    this->declare_parameter("use_kalman_filter", true);
    this->declare_parameter("use_calibration", false);
    //卡尔曼滤波参数声明
    this->declare_parameter("process_noise", 0.01);
    this->declare_parameter("gyro_noise", 0.002);
    this->declare_parameter("accel_noise", 0.01);
    this->declare_parameter("static_process_noise", 0.005);
    this->declare_parameter("dynamic_process_noise", 0.1);

    //获取参数
    this->get_parameter("input_imu_topic", input_imu_topic_);
    this->get_parameter("output_imu_topic", output_imu_topic_);
    this->get_parameter("imu_frame", imu_frame_);
    this->get_parameter("base_frame", base_frame_);
    this->get_parameter("publish_tf", publish_tf_);
    this->get_parameter("imu_position_x", imu_position_x_);
    this->get_parameter("imu_position_y", imu_position_y_);
    this->get_parameter("imu_position_z", imu_position_z_);
    this->get_parameter("imu_rotation_x", imu_rotation_x_);
    this->get_parameter("imu_rotation_y", imu_rotation_y_);
    this->get_parameter("imu_rotation_z", imu_rotation_z_);
    this->get_parameter("imu_rotation_w", imu_rotation_w_);
    this->get_parameter("use_kalman_filter", use_kalman_filter_);
    this->get_parameter("use_calibration", use_calibration_);
    this->get_parameter("process_noise", process_noise_);
    this->get_parameter("gyro_noise", gyro_noise_);
    this->get_parameter("accel_noise", accel_noise_);
    this->get_parameter("static_process_noise", static_process_noise_);
    this->get_parameter("dynamic_process_noise", dynamic_process_noise_);

    /* //初始化旋转矩阵
    initialize_rotation_matrix(); */

    //初始化卡尔曼滤波器
    if(use_kalman_filter_) {
        angular_vel_x_filter_ = std::make_unique<SimpleKalmanFilter>(process_noise_, gyro_noise_);
        angular_vel_y_filter_ = std::make_unique<SimpleKalmanFilter>(process_noise_, gyro_noise_);
        angular_vel_z_filter_ = std::make_unique<SimpleKalmanFilter>(process_noise_, gyro_noise_);
        linear_acc_x_filter_ = std::make_unique<SimpleKalmanFilter>(process_noise_, accel_noise_);
        linear_acc_y_filter_ = std::make_unique<SimpleKalmanFilter>(process_noise_, accel_noise_);
        linear_acc_z_filter_ = std::make_unique<SimpleKalmanFilter>(process_noise_, accel_noise_);      
    }

    //发布静态tf变换
    if(publish_tf_) {
        publish_static_tf();
    }

    //创建订阅者和发布者
    imu_sub_ = this->create_subscription<sensor_msgs::msg::Imu>(input_imu_topic_, 10,
                std::bind(&ImuProcessor::imu_callback, this, std::placeholders::_1));
    imu_pub_ = this->create_publisher<sensor_msgs::msg::Imu>(output_imu_topic_, 10);

    //创建校准服务
    calibrate_service_ = this->create_service<std_srvs::srv::Trigger>(
        "calibrate_imu",
        std::bind(&ImuProcessor::calibrate_callback, this, 
                 std::placeholders::_1, std::placeholders::_2));

    RCLCPP_INFO(this->get_logger(), "IMU Filter Node started");
    RCLCPP_INFO(this->get_logger(), "Subscribing to: %s", input_imu_topic_.c_str());
    RCLCPP_INFO(this->get_logger(), "Publishing to: %s", output_imu_topic_.c_str());
    if(use_kalman_filter_) {
        RCLCPP_INFO(this->get_logger(), "Using kalman filter with process_noise: %f, gyro_noise: %f, accel_noise: %f", 
                   process_noise_, gyro_noise_, accel_noise_);
    }
    if(use_calibration_) {
        RCLCPP_INFO(this->get_logger(), "IMU calibration enabled");
    }
}


void ImuProcessor::publish_static_tf() {
    geometry_msgs::msg::TransformStamped static_transform;
    static_transform.header.stamp = this->now();
    static_transform.header.frame_id = base_frame_;
    static_transform.child_frame_id = imu_frame_;

    //设置IMU安装位置
    static_transform.transform.translation.x = imu_position_x_;
    static_transform.transform.translation.y = imu_position_y_;
    static_transform.transform.translation.z = imu_position_z_;

    //设置IMU旋转
    static_transform.transform.rotation.x = imu_rotation_x_;
    static_transform.transform.rotation.y = imu_rotation_y_;
    static_transform.transform.rotation.z = imu_rotation_z_;
    static_transform.transform.rotation.w = imu_rotation_w_;

    static_broadcaster_ = std::make_shared<tf2_ros::StaticTransformBroadcaster>(this);
    static_broadcaster_->sendTransform(static_transform);

    RCLCPP_INFO(this->get_logger(), "Published static TF: %s -> %s",
                base_frame_.c_str(), imu_frame_.c_str());
        
}

bool ImuProcessor::is_zero_covariance(const std::array<double, 9>& cov) {
    for(double c : cov) {
        if(c != 0.0) return false;
    }
    return true;
}

bool ImuProcessor::is_vehicle_stationary(const sensor_msgs::msg::Imu& imu_data) {
    double angular_vel_threshold = 0.02;  // rad/s
    double linear_acc_threshold = 0.2;    // m/s²
    
    double angular_vel_mag = get_angular_velocity_magnitude(imu_data);
    double linear_acc_mag = get_linear_acceleration_magnitude(imu_data);
    
    return (angular_vel_mag < angular_vel_threshold) && 
           (linear_acc_mag < linear_acc_threshold);
}

double ImuProcessor::get_angular_velocity_magnitude(const sensor_msgs::msg::Imu& imu_data) {
    return sqrt(
        imu_data.angular_velocity.x * imu_data.angular_velocity.x +
        imu_data.angular_velocity.y * imu_data.angular_velocity.y +
        imu_data.angular_velocity.z * imu_data.angular_velocity.z
    );
}

double ImuProcessor::get_linear_acceleration_magnitude(const sensor_msgs::msg::Imu& imu_data) {
    return sqrt(
        imu_data.linear_acceleration.x * imu_data.linear_acceleration.x +
        imu_data.linear_acceleration.y * imu_data.linear_acceleration.y +
        imu_data.linear_acceleration.z * imu_data.linear_acceleration.z
    );
}


void ImuProcessor::imu_callback(const sensor_msgs::msg::Imu::SharedPtr raw_imu) {
    //创建矫正后的imu消息
    sensor_msgs::msg::Imu corrected_imu = *raw_imu;

    //更新坐标系
    corrected_imu.header.frame_id = imu_frame_;

    //设置合理的协方差
    if(is_zero_covariance(raw_imu->orientation_covariance)) {
        corrected_imu.orientation_covariance = {{0.05, 0, 0, 0, 0.05, 0, 0, 0, 0.05}};
    }
    if(is_zero_covariance(raw_imu->angular_velocity_covariance)) {
        corrected_imu.angular_velocity_covariance = {{0.002, 0, 0, 0, 0.002, 0, 0, 0, 0.002}};
    }
    if(is_zero_covariance(raw_imu->linear_acceleration_covariance)) {
        corrected_imu.linear_acceleration_covariance = {{0.01, 0, 0, 0, 0.01, 0, 0, 0, 0.005}};
    }

    

    //应用坐标变换（针对激光雷达倒置情况）
    //如果激光雷达倒置，需要对IMU数据进行相应变换
    corrected_imu.angular_velocity.y = -corrected_imu.angular_velocity.y;
    corrected_imu.angular_velocity.z = -corrected_imu.angular_velocity.z;
    /* corrected_imu.linear_acceleration.y = -corrected_imu.linear_acceleration.y; */
    /* corrected_imu.linear_acceleration.z = -corrected_imu.linear_acceleration.z; */


     //应用零偏校准
    if(use_calibration_ && calibrator_.get_sample_count() > 0) {
        corrected_imu = calibrator_.apply_calibration(corrected_imu);
    }

    //更新校准采样
    if(calibrator_.is_calibrating()) {
        calibrator_.update_calibration(raw_imu);
        if(calibrator_.get_sample_count() >= calibrator_.get_max_samples()) {
            calibrator_.compute_bias();
            calibrator_.start_calibration(); // 停止校准
            RCLCPP_INFO(this->get_logger(), "IMU calibration completed");
        } else {
            RCLCPP_INFO_THROTTLE(this->get_logger(), *this->get_clock(), 1000,
                "IMU calibration progress: %zu/%zu samples", 
                calibrator_.get_sample_count(), calibrator_.get_max_samples());
        }
    }

    //应用卡尔曼滤波
    if(use_kalman_filter_) {

         //根据运动状态调整过程噪声
        double current_process_noise = is_vehicle_stationary(corrected_imu) ? 
                                      static_process_noise_ : dynamic_process_noise_;

        //动态更新滤波器的过程噪声
        angular_vel_x_filter_->update_process_noise(current_process_noise);
        angular_vel_y_filter_->update_process_noise(current_process_noise);
        angular_vel_z_filter_->update_process_noise(current_process_noise);
        linear_acc_x_filter_->update_process_noise(current_process_noise);
        linear_acc_y_filter_->update_process_noise(current_process_noise);
        linear_acc_z_filter_->update_process_noise(current_process_noise);

        corrected_imu.angular_velocity.x = angular_vel_x_filter_->filter(corrected_imu.angular_velocity.x);
        corrected_imu.angular_velocity.y = angular_vel_y_filter_->filter(corrected_imu.angular_velocity.y);
        corrected_imu.angular_velocity.z = angular_vel_z_filter_->filter(corrected_imu.angular_velocity.z);
        corrected_imu.linear_acceleration.x = linear_acc_x_filter_->filter(corrected_imu.linear_acceleration.x);
        corrected_imu.linear_acceleration.y = linear_acc_y_filter_->filter(corrected_imu.linear_acceleration.y);
        corrected_imu.linear_acceleration.z = linear_acc_z_filter_->filter(corrected_imu.linear_acceleration.z);
    }
    
    imu_pub_->publish(corrected_imu);
}

void ImuProcessor::calibrate_callback(
    const std::shared_ptr<std_srvs::srv::Trigger::Request> ,
    std::shared_ptr<std_srvs::srv::Trigger::Response> response) {
    
    calibrator_.start_calibration();
    response->success = true;
    response->message = "开始IMU零偏校准, 请保持IMU静止...";
    RCLCPP_INFO(this->get_logger(), "IMU calibration started. Please keep IMU stationary.");
}


int main(int argc, char** argv) 
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<ImuProcessor>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}